'''
####
Date : 2025-1-9
Notes: UI界面显示
此脚本主要功能:
1. 获取视觉图像

####
'''

# -*- coding: utf-8 -*-
import cv2
import os
import logging
import time
import queue
import threading
import imutils
import numpy as np
from paddleocr import PaddleOCR

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)


current_dir = os.path.dirname(os.path.abspath(__file__))
res_file_path = os.path.join(current_dir, "..", "img/res", "result.png")

kernel_3 = np.ones((3, 3), np.uint8)  # 3x3的卷积核
kernel_5 = np.ones((5, 5), np.uint8)  
kernel_7 = np.ones((7, 7), np.uint8)

draw_color = (0,255,0)

Lower_red = np.array([0, 97, 113])  
Upper_red = np.array([180, 255, 255])
draw_red = draw_color
red = [Lower_red, Upper_red, 'red',draw_red]

Lower_blue = np.array([106, 219, 142])
Upper_blue = np.array([111, 255, 225])
draw_blue =draw_color
blue = [Lower_blue, Upper_blue, 'blue',draw_blue]

Lower_green = np.array([39, 60, 133])
Upper_green = np.array([96, 255, 255])
draw_green =draw_color
green = [Lower_green, Upper_green, 'green',draw_green]

Lower_yellow = np.array([21, 80, 200])   
Upper_yellow = np.array([37, 99, 255])
draw_yellow = draw_color
yellow = [Lower_yellow, Upper_yellow, 'yellow',draw_yellow]

Lower_black = np.array([100, 118, 37])
Upper_black = np.array([107, 193, 69])
draw_black = (0,0,0)
black = [Lower_black, Upper_black, 'black',draw_black]

color_list = [red,blue,green,yellow]


### version 1.0.2
class Camera_Thread(threading.Thread):
    
    def __init__(self, label):
        super().__init__()
        self.stop = True
        self.Is_exit = False
        self.cap = None
        ## 存放原始图像帧
        self.raw_frame_queue = queue.Queue(maxsize=1)  
        ## 存放处理后的图像帧
        self.processed_frame_queue = queue.Queue(maxsize=1)  
        self.label = label
        # self.label_res = label_res
        ## 控制是否进行二维码检测的标志
        self.detect_qr_code = False  
        self.resultInfo = None
        ## 视觉检测参数配置
        # self.parameterInfo = list()
        self.parameterInfo = []
        self.seen_items = set()

        ## 可以根据需要更改语言参数
        self.ocr = PaddleOCR(use_angle_cls=True, lang='ch')  

    def exit(self):
        self.stop = True
        self.Is_exit = True
        if self.cap is not None:
            self.cap.release()

    def add_parameter_info(self, msg):
            # 如果msg不在seen_items中，则添加到parameterInfo和seen_items
            if msg not in self.seen_items:
                self.parameterInfo.append(msg)
                self.seen_items.add(msg)

    ### 二维码识别--调用微信开源库
    def QR_test(self, img):
        try:
            detect_obj = cv2.wechat_qrcode_WeChatQRCode('detect.prototxt', 'detect.caffemodel', 'sr.prototxt', 'sr.caffemodel')
            res, points = detect_obj.detectAndDecode(img)
            print(f"QR_test res is {res}")
            
            # if res:
            #     self.parameterInfo.append(f"检测识别结果: {res}")
            # else:
            #     self.parameterInfo.append("未检测到二维码或无法解码")
            
            if points is not None:
                for idx, qr_points in enumerate(points):
                    ## 绘制中心点和边界线
                    x0, y0 = qr_points[0]
                    x1, y1 = qr_points[2]
                    center_x = (x0 + x1) / 2
                    center_y = (y0 + y1) / 2
                    center_point = (int(center_x), int(center_y))
                    cv2.circle(img, center_point, 5, (0, 255, 0), -1)

                    color = (0, 0, 255)
                    thick = 3
                    for p in [(0, 1), (1, 2), (2, 3), (3, 0)]:
                        start = tuple(map(int, qr_points[p[0]]))
                        end = tuple(map(int, qr_points[p[1]]))
                        cv2.line(img, start, end, color, thick)

                    ## 显示解码结果
                    qr_code_text = res[idx]
                    cv2.putText(img, qr_code_text[:1], (int(x0), int(y0) - 10), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 125), 2)
                    
                    ## 添加打印日志
                    self.resultInfo = "QR: " + qr_code_text
                    
                    detect_model_info = "定义存储的Caffe模型: \n detect.prototxt, detect.caffemodel"
                    sr_model_info = "超分辨率模型: \n sr.prototxt, sr.caffemodel"

                    # 尝试添加信息，如果不存在则添加
                    self.add_parameter_info(detect_model_info)
                    self.add_parameter_info(sr_model_info)
            
            return img
        
        except Exception as e:
            print(f"Error in QR_test: {e}")
            return img

    ### OCR识别--百度飞桨OCRv4 
    def OCR_test(self, img):
        try:
            # 执行文字检测和识别
            result = self.ocr.ocr(img, cls=True)
            print(f"ocr result is {result}")

            # 如果有识别结果，则处理并显示
            if result is not None:
                for res in result:
                    for line in res:
                        print(f"line {line}")
                        boxes = line[0]  # 文本框坐标
                        txts, scores = line[1]


                        # pts = [tuple(map(int, point)) for point in boxes]
                        # cv2.polylines(img, [np.array(pts)], isClosed=True, color=(0, 255, 0), thickness=2)
                        color = (0, 0, 255)
                        thick = 3
                        for p in [(0, 1), (1, 2), (2, 3), (3, 0)]:
                            start = tuple(map(int, boxes[p[0]]))
                            end = tuple(map(int, boxes[p[1]]))
                            cv2.line(img, start, end, color, thick)
                        
                        # 在文本框上方显示识别出的文本
                        text_x, text_y = int(boxes[0][0]), int(boxes[1][1]) - 10
                        # print(text_x, text_y)
                        # print(f"文本Text: {txts}, Confidence: {scores}")
                        
                        cv2.putText(img, txts, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)
                        # cv2.putText(img, qr_code_text[:1], (int(x0), int(y0) - 10), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 125), 2)

                        # font = cv2.FONT_HERSHEY_SIMPLEX
                        # cv2.putText(img, txts, (int(boxes[0][0]), int(boxes[0][1])-10), font, 1, (0, 0, 255), 2, cv2.LINE_AA)
                        
                    # 更新结果信息
                    self.resultInfo = "文本: " + txts
                
                # 添加模型信息（这里假设你想要记录一些关于模型的信息）
                model_info = "使用的PaddleOCR模型: \n ch_PP-OCRv4_det_infer, ch_PP-OCRv4_rec_infer \n"
                self.add_parameter_info(model_info)
                
            return img
            
        except Exception as e:
            print(f"Error in OCR_test: {e}")
            return img

    ### 形状识别--OpenCV
    def shape_test(self, frame):
        max_area = 200000
        min_area = 1800
        imgGrey = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        _, thresh = cv2.threshold(imgGrey, 80, 255, cv2.THRESH_BINARY)
        
        # cv2.imwrite("./img/res/thresh.png", thresh)
        inverted_binary = cv2.bitwise_not(thresh)
        # cv2.imwrite("./img/res/inverted_binary.png", inverted_binary)
        contours, _ = cv2.findContours(inverted_binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
        
        for contour in contours:
            contour_area = cv2.contourArea(contour)
            if contour_area > min_area and contour_area < max_area:
                approx = cv2.approxPolyDP(contour, 0.01* cv2.arcLength(contour, True), True)
                cv2.drawContours(frame, [approx], 0, (0,0,255), 5)
                x = approx.ravel()[0]
                y = approx.ravel()[1] - 5
                if len(approx) == 3:
                    cv2.putText(frame, "Triangle", (x, y), cv2.FONT_HERSHEY_COMPLEX, 1, (0,0,255))
                    self.resultInfo = f"形状 三角形"

                elif len(approx) == 4:
                    x1, y1, w, h = cv2.boundingRect(approx)
                    aspectRatio = float(w)/h
                    
                    if aspectRatio >= 0.95 and aspectRatio <= 1.05:
                        cv2.putText(frame, "Square", (x, y), cv2.FONT_HERSHEY_COMPLEX, 1, (0,0,255))
                        self.resultInfo = f"形状 正方形"

                    else:
                        cv2.putText(frame, "Rectangle", (x, y), cv2.FONT_HERSHEY_COMPLEX, 1, (0,0,255))
                        self.resultInfo = f"形状 矩形"
                
                elif len(approx) == 5:
                    cv2.putText(frame, "Pentagon", (x, y), cv2.FONT_HERSHEY_COMPLEX, 1, (0,0,255))
                    self.resultInfo = f"形状 五边形"

                elif len(approx) == 6:
                    cv2.putText(frame, "Hexagon", (x, y), cv2.FONT_HERSHEY_COMPLEX, 1, (0,0,255))
                    self.resultInfo = f"形状 六边形"
                
                elif len(approx) == 10:
                    cv2.putText(frame, "Star", (x, y), cv2.FONT_HERSHEY_COMPLEX, 1, (0,0,255))
                    self.resultInfo = f"形状 五角星"
                
                else:
                    cv2.putText(frame, "Circle", (x, y), cv2.FONT_HERSHEY_COMPLEX, 1, (0,0,255))
                    self.resultInfo = f"形状 圆形"
                # 添加参数信息
                model_info = f" 二值化阈值为 [80, 255], 检测区域面积[1800, 200000]"
                self.add_parameter_info(model_info)
        
        cv2.imwrite(res_file_path, frame) # 移到循环外，确保所有检测结果都绘制在图像上后再保存
        return frame

    ### 检测识别-- 颜色识别
    def color_test(self, image):
        w = 70
        h = 70
        ##### 把BGR图像转换为HSV格式
        HSV = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        kernel = np.ones((3,3), np.uint8) # 定义卷积核大小
        
        for color in color_list:
            mask_red = cv2.inRange(HSV, color[0], color[1])
            
            ### 形态学去噪
            mask_red = cv2.morphologyEx(mask_red, cv2.MORPH_OPEN, kernel) # 开运算
            
            ### 二值化
            ret, binary = cv2.threshold(mask_red, 0, 255, cv2.THRESH_BINARY)
            cv2.imwrite("./img/res/binary.png", binary)
            cnts = cv2.findContours(binary.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
            cnts = imutils.grab_contours(cnts)

            if cnts:
                for c in cnts:
                    area = cv2.contourArea(c) # 计算轮廓面积
                    if area < 1800 or area > 20000: # 面积筛选
                        continue
                    
                    M0_red = cv2.moments(c)
                    if M0_red["m00"] != 0: # 确保避免除以零错误
                        car_video_location = [int(M0_red["m10"] / M0_red["m00"]), int(M0_red["m01"] / M0_red["m00"])]
                        ptLeftTop = (car_video_location[0] - h, car_video_location[1] - w)
                        ptRightBottom = (car_video_location[0] + h, car_video_location[1] + w)
                        point_color = color[3]
                        thickness = 2 
                        lineType = 4
                        
                        cv2.rectangle(image, ptLeftTop, ptRightBottom, point_color, thickness, lineType)
                        cv2.putText(image, color[2], (car_video_location[0]-20, car_video_location[1]-30),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color[3], 2)
                        # 更新结果信息
                        self.resultInfo = f"颜色: {color[2]}"
        
                        # 添加参数信息
                        model_info = f" {color[2]} 匹配 \n HSV Min: {color[0]} \n HSV Max: {color[1]}"
                        self.add_parameter_info(model_info)
                
        cv2.imwrite(res_file_path, image)
        
        return image

    def run(self):
        try:
            current_dir = os.path.dirname(os.path.abspath(__file__))
            self.cap = cv2.VideoCapture(18, cv2.CAP_V4L2)
            self.cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'))
            ori_file_path = os.path.join(current_dir, "..", "img/ori", "frame.png")

            while not self.Is_exit:
                if not self.stop:
                    ret, frame = self.cap.read()
                    if ret:
                        image = frame.copy()
                        cv2.imwrite(ori_file_path, image)
                        
                        ### 将原始图像放入原始图像队列
                        try:
                            if not self.raw_frame_queue.full():
                                self.raw_frame_queue.put(frame)
                            else:
                                self.raw_frame_queue.get_nowait()  
                                self.raw_frame_queue.put(frame) 
                        except queue.Empty:
                            pass

                        ### (自动模式)如果启用了二维码检测，则处理图像并放入处理图像队列
                        if self.detect_qr_code:
                            processed_image = self.QR_test(image)
                            try:
                                if not self.raw_frame_queue.full():
                                    self.raw_frame_queue.put(processed_image)
                                else:
                                    self.raw_frame_queue.get_nowait() 
                                    self.raw_frame_queue.put(processed_image) 
                            except queue.Empty:
                                pass

                else:
                    ### (手动模式) 如果启用了二维码检测，则处理图像并放入处理图像队列，处理img/ori下保存的图像
                    ### 下面手动选择视觉检测方法
                    if self.detect_qr_code:
                        image = cv2.imread(ori_file_path)
                        ## 二维码识别
                        # processed_image = self.QR_test(image)
                        ## OCR识别
                        # processed_image = self.OCR_test(image)
                        ## 形状识别
                        processed_image = self.shape_test(image)
                        ## 颜色识别
                        # processed_image = self.color_test(image)
                        
                        try:
                            if not self.raw_frame_queue.full():
                                self.raw_frame_queue.put(processed_image)
                            else:
                                self.raw_frame_queue.get_nowait() 
                                self.raw_frame_queue.put(processed_image) 
                        except queue.Empty:
                            pass

                        self.stop = True
                        self.detect_qr_code = False
                    time.sleep(0.1)
                    

        except Exception as e:
            print(f"Error in run: {e}")
